"""
Cartesian grid implementations.
"""

using LinearAlgebra
using StaticArrays

"""
    UnitGrid{D} <: AbstractCartesianGrid

Simple equidistant grid on the unit hypercube [0,1]^D.

# Fields
- `dims::NTuple{D,Int}`: Number of grid points in each dimension

# Examples
```julia
# 1D grid with 64 points
grid = UnitGrid([64])

# 2D grid with 64x32 points  
grid = UnitGrid([64, 32])
```
"""
struct UnitGrid{D} <: AbstractCartesianGrid
    dims::NTuple{D,Int}
    
    function UnitGrid{D}(dims::NTuple{D,Int}) where D
        all(d -> d > 0, dims) || throw(ArgumentError("Grid dimensions must be positive"))
        new{D}(dims)
    end
end

UnitGrid(dims::AbstractVector{Int}) = UnitGrid{length(dims)}(Tuple(dims))
UnitGrid(dims::Vararg{Int}) = UnitGrid{length(dims)}(dims)

# Interface implementation for UnitGrid
ndim(::UnitGrid{D}) where D = D
Base.size(grid::UnitGrid) = grid.dims

function extent(::UnitGrid{D}) where D
    ntuple(_ -> (0.0, 1.0), D)
end

function spacing(grid::UnitGrid)
    ntuple(i -> 1.0 / (grid.dims[i] - 1), ndim(grid))
end

function coordinate(grid::UnitGrid, idx)
    coords = ntuple(ndim(grid)) do i
        (idx[i] - 1) / (grid.dims[i] - 1)
    end
    return SVector(coords)
end

function index(grid::UnitGrid, coord)
    ntuple(ndim(grid)) do i
        round(Int, coord[i] * (grid.dims[i] - 1) + 1)
    end
end

"""
    CartesianGrid{D} <: AbstractCartesianGrid

General Cartesian grid with arbitrary rectangular domain.

# Fields
- `dims::NTuple{D,Int}`: Number of grid points in each dimension
- `bounds::NTuple{D,Tuple{Float64,Float64}}`: Physical bounds [min, max] for each dimension

# Examples
```julia
# 1D grid from 0 to 10 with 100 points
grid = CartesianGrid([100], [(0.0, 10.0)])

# 2D grid with different bounds and resolution
grid = CartesianGrid([64, 32], [(0.0, 5.0), (-1.0, 1.0)])
```
"""
struct CartesianGrid{D} <: AbstractCartesianGrid
    dims::NTuple{D,Int}
    bounds::NTuple{D,Tuple{Float64,Float64}}
    
    function CartesianGrid{D}(dims::NTuple{D,Int}, bounds::NTuple{D,Tuple{Float64,Float64}}) where D
        all(d -> d > 0, dims) || throw(ArgumentError("Grid dimensions must be positive"))
        all(b -> first(b) < last(b), bounds) || throw(ArgumentError("Invalid bounds"))
        new{D}(dims, bounds)
    end
end

function CartesianGrid(dims::AbstractVector{Int}, bounds::AbstractVector{Tuple{Float64,Float64}})
    D = length(dims)
    length(bounds) == D || throw(ArgumentError("Dimensions and bounds must have same length"))
    CartesianGrid{D}(Tuple(dims), Tuple(bounds))
end

function CartesianGrid(dims::AbstractVector{Int}, bounds::AbstractVector{AbstractVector})
    bounds_tuples = [Tuple(b) for b in bounds]
    CartesianGrid(dims, bounds_tuples)
end

# Interface implementation for CartesianGrid
ndim(::CartesianGrid{D}) where D = D
Base.size(grid::CartesianGrid) = grid.dims
extent(grid::CartesianGrid) = grid.bounds

function spacing(grid::CartesianGrid)
    ntuple(ndim(grid)) do i
        bounds = grid.bounds[i]
        (last(bounds) - first(bounds)) / (grid.dims[i] - 1)
    end
end

function coordinate(grid::CartesianGrid, idx)
    coords = ntuple(ndim(grid)) do i
        bounds = grid.bounds[i]
        first(bounds) + (idx[i] - 1) * spacing(grid)[i]
    end
    return SVector(coords)
end

function index(grid::CartesianGrid, coord)
    ntuple(ndim(grid)) do i
        bounds = grid.bounds[i]
        s = spacing(grid)[i]
        round(Int, (coord[i] - first(bounds)) / s + 1)
    end
end

# Convenience constructors
"""
    CartesianGrid(bounds, num_points)

Create Cartesian grid with uniform resolution.

# Arguments
- `bounds`: Vector of [min, max] pairs for each dimension
- `num_points`: Number of points (same for all dimensions)

# Examples
```julia
grid = CartesianGrid([[0, 1], [0, 1]], 64)  # 64x64 grid on unit square
```
"""
function CartesianGrid(bounds::AbstractVector, num_points::Int)
    dims = fill(num_points, length(bounds))
    CartesianGrid(dims, bounds)
end

"""
    get_line(grid::CartesianGrid, start, end, num_points=100)

Extract a line through the grid from start to end coordinates.
"""
function get_line(grid::CartesianGrid, start, finish, num_points=100)
    t_vals = range(0, 1, length=num_points)
    points = [start .+ t .* (finish .- start) for t in t_vals]
    return points
end

"""
    get_image_data(grid::CartesianGrid{2}, field_data)

Prepare 2D field data for plotting/visualization.
"""
function get_image_data(grid::CartesianGrid{2}, field_data)
    x_axis = get_axis(grid, 1)
    y_axis = get_axis(grid, 2)
    return x_axis, y_axis, field_data'
end

# Grid refinement operations
"""
    refine(grid::CartesianGrid, factor::Int=2)

Create a refined grid with `factor` times more points in each dimension.
"""
function refine(grid::CartesianGrid{D}, factor::Int=2) where D
    new_dims = ntuple(i -> (grid.dims[i] - 1) * factor + 1, D)
    CartesianGrid{D}(new_dims, grid.bounds)
end

"""
    coarsen(grid::CartesianGrid, factor::Int=2)

Create a coarsened grid with `factor` times fewer points in each dimension.
"""
function coarsen(grid::CartesianGrid{D}, factor::Int=2) where D
    new_dims = ntuple(i -> div(grid.dims[i] - 1, factor) + 1, D)
    CartesianGrid{D}(new_dims, grid.bounds)
end

# Advanced coordinate transformations and indexing
"""
    coordinate_range(grid::AbstractCartesianGrid, axis::Int) -> Vector

Get all coordinate values along a specific axis.
"""
function coordinate_range(grid::AbstractCartesianGrid, axis::Int)
    bounds = extent(grid)[axis]
    n = size(grid)[axis]
    return collect(range(first(bounds), last(bounds), length=n))
end

"""
    cell_coordinates(grid::AbstractCartesianGrid) -> Tuple

Get the coordinate arrays for all axes (cell-centered coordinates).
"""
function cell_coordinates(grid::G) where {G <: AbstractCartesianGrid}
    D = ndim(grid)
    ntuple(i -> coordinate_range(grid, i), D)
end

"""
    grid_coordinates(grid::AbstractCartesianGrid) -> Tuple

Get the coordinate arrays for all axes (for meshing).
"""
function grid_coordinates(grid::G) where {G <: AbstractCartesianGrid}
    D = ndim(grid)
    axes = cell_coordinates(grid)
    if D == 1
        return (axes[1],)
    elseif D == 2
        x, y = axes
        X = repeat(x, 1, length(y))
        Y = repeat(y', length(x), 1)
        return (X, Y)
    elseif D == 3
        x, y, z = axes
        nx, ny, nz = length(x), length(y), length(z)
        X = repeat(reshape(x, nx, 1, 1), 1, ny, nz)
        Y = repeat(reshape(y, 1, ny, 1), nx, 1, nz)
        Z = repeat(reshape(z, 1, 1, nz), nx, ny, 1)
        return (X, Y, Z)
    else
        throw(ArgumentError("Grid coordinates only supported for D ≤ 3"))
    end
end

"""
    contains_point(grid::AbstractCartesianGrid, point) -> Bool

Check if a point is within the grid domain.
"""
function contains_point(grid::AbstractCartesianGrid, point)
    ext = extent(grid)
    for (i, coord) in enumerate(point)
        bounds = ext[i]
        if coord < first(bounds) || coord > last(bounds)
            return false
        end
    end
    return true
end

"""
    closest_index(grid::AbstractCartesianGrid, point) -> Tuple

Find the closest grid index to a given point, clamped to grid bounds.
"""
function closest_index(grid::AbstractCartesianGrid, point)
    idx = index(grid, point)
    dims = size(grid)
    return ntuple(i -> clamp(idx[i], 1, dims[i]), ndim(grid))
end

# Ghost cell and boundary handling
"""
    get_boundary_indices(grid::AbstractCartesianGrid, axis::Int, side::Symbol) -> Vector{Int}

Get the indices of boundary points along a specific axis.
- side: :lower for minimum boundary, :upper for maximum boundary
"""
function get_boundary_indices(grid::G, axis::Int, side::Symbol) where {G <: AbstractCartesianGrid}
    dims = size(grid)
    if side == :lower
        return [1]
    elseif side == :upper
        return [dims[axis]]
    else
        throw(ArgumentError("side must be :lower or :upper"))
    end
end

"""
    get_ghost_cell_indices(grid::AbstractCartesianGrid, axis::Int, side::Symbol, num_ghost::Int=1) -> Vector{Int}

Get indices for ghost cells outside the domain boundaries.
"""
function get_ghost_cell_indices(grid::AbstractCartesianGrid, axis::Int, side::Symbol, num_ghost::Int=1)
    dims = size(grid)
    if side == :lower
        return collect(1-num_ghost:0)
    elseif side == :upper
        return collect(dims[axis]+1:dims[axis]+num_ghost)
    else
        throw(ArgumentError("side must be :lower or :upper"))
    end
end

"""
    add_ghost_cells(grid::AbstractCartesianGrid{D}, num_ghost::Int=1) -> CartesianGrid{D}

Create a new grid with ghost cells added on all boundaries.
"""
function add_ghost_cells(grid::G, num_ghost::Int=1) where {G <: AbstractCartesianGrid}
    D = ndim(grid)
    old_dims = size(grid)
    new_dims = ntuple(i -> old_dims[i] + 2*num_ghost, D)
    
    # Extend bounds to include ghost cells
    old_bounds = extent(grid)
    spacings = spacing(grid)
    new_bounds = ntuple(i -> begin
        old_bound = old_bounds[i]
        dx = spacings[i]
        (first(old_bound) - num_ghost*dx, last(old_bound) + num_ghost*dx)
    end, D)
    
    if isa(grid, UnitGrid)
        # Can't add ghost cells to UnitGrid directly, convert to CartesianGrid
        return CartesianGrid(collect(new_dims), collect(new_bounds))
    else
        return CartesianGrid(collect(new_dims), collect(new_bounds))
    end
end

# Integration and interpolation
"""
    integrate_simpson(grid::UnitGrid{1}, data::Vector) -> Float64

Integrate 1D data using Simpson's rule.
"""
function integrate_simpson(grid::UnitGrid{1}, data::Vector)
    n = length(data)
    if n != size(grid)[1]
        throw(DimensionMismatch("Data length must match grid size"))
    end
    
    if n < 3
        throw(ArgumentError("Simpson's rule requires at least 3 points"))
    end
    
    dx = spacing(grid)[1]
    if isodd(n)
        # Standard Simpson's rule
        integral = data[1] + data[n]
        for i in 2:2:n-1
            integral += 4 * data[i]
        end
        for i in 3:2:n-2
            integral += 2 * data[i]
        end
        return integral * dx / 3
    else
        # Simpson's 3/8 rule for the first interval, then standard
        integral = (3/8) * dx * (data[1] + 3*data[2] + 3*data[3] + data[4])
        for i in 5:2:n-1
            integral += (dx/3) * (data[i-2] + 4*data[i-1] + data[i])
        end
        return integral
    end
end

"""
    integrate_trapezoidal(grid::AbstractCartesianGrid, data::Array) -> Float64

Integrate multi-dimensional data using trapezoidal rule.
"""
function integrate_trapezoidal(grid::G, data::Array{T,D}) where {D,T,G <: AbstractCartesianGrid}
    if size(data) != size(grid)
        throw(DimensionMismatch("Data size must match grid size"))
    end
    
    spacings = spacing(grid)
    cell_vol = prod(spacings)
    
    return sum(data) * cell_vol
end

"""
    interpolate_linear(grid::AbstractCartesianGrid, data::Array, point) -> Float64

Linear interpolation at a given point.
"""
function interpolate_linear(grid::G, data::Array{T,D}, point) where {D,T,G <: AbstractCartesianGrid}
    if size(data) != size(grid)
        throw(DimensionMismatch("Data size must match grid size"))
    end
    
    if !contains_point(grid, point)
        throw(BoundsError("Point $point is outside grid domain"))
    end
    
    # Find surrounding grid points
    dims = size(grid)
    ext = extent(grid)
    spacings = spacing(grid)
    
    # Convert point to grid coordinates
    grid_coords = ntuple(D) do i
        bounds = ext[i]
        (point[i] - first(bounds)) / spacings[i] + 1
    end
    
    # Find integer indices and fractional parts
    indices = ntuple(i -> floor(Int, grid_coords[i]), D)
    fractions = ntuple(i -> grid_coords[i] - indices[i], D)
    
    # Clamp indices to valid range
    indices = ntuple(i -> clamp(indices[i], 1, dims[i]-1), D)
    
    if D == 1
        i = indices[1]
        α = fractions[1]
        return (1-α) * data[i] + α * data[i+1]
    elseif D == 2
        i, j = indices
        α, β = fractions
        return ((1-α) * (1-β) * data[i, j] + 
                α * (1-β) * data[i+1, j] + 
                (1-α) * β * data[i, j+1] + 
                α * β * data[i+1, j+1])
    elseif D == 3
        i, j, k = indices
        α, β, γ = fractions
        return ((1-α) * (1-β) * (1-γ) * data[i, j, k] +
                α * (1-β) * (1-γ) * data[i+1, j, k] +
                (1-α) * β * (1-γ) * data[i, j+1, k] +
                (1-α) * (1-β) * γ * data[i, j, k+1] +
                α * β * (1-γ) * data[i+1, j+1, k] +
                α * (1-β) * γ * data[i+1, j, k+1] +
                (1-α) * β * γ * data[i, j+1, k+1] +
                α * β * γ * data[i+1, j+1, k+1])
    else
        throw(ArgumentError("Linear interpolation only supported for D ≤ 3"))
    end
end

# Data extraction methods
"""
    get_line_data(grid::AbstractCartesianGrid, data::Array, axis::Int; index::Int=0) -> Tuple

Extract a line of data along a specific axis.
If index=0, extracts through the center of other dimensions.
"""
function get_line_data(grid::G, data::Array{T,D}, axis::Int; index::Int=0) where {D,T,G <: AbstractCartesianGrid}
    if size(data) != size(grid)
        throw(DimensionMismatch("Data size must match grid size"))
    end
    
    dims = size(grid)
    
    # Determine indices for other dimensions
    other_indices = ntuple(D) do i
        if i == axis
            Colon()
        else
            index == 0 ? dims[i] ÷ 2 + 1 : index
        end
    end
    
    line_data = data[other_indices...]
    coords = coordinate_range(grid, axis)
    
    return coords, line_data
end

"""
    get_cross_section(grid::CartesianGrid{3}, data::Array{T,3}, axis::Int; index::Int=0) -> Tuple

Extract a 2D cross-section from 3D data perpendicular to the specified axis.
"""
function get_cross_section(grid::CartesianGrid{3}, data::Array{T,3}, axis::Int; index::Int=0) where T
    if size(data) != size(grid)
        throw(DimensionMismatch("Data size must match grid size"))
    end
    
    dims = size(grid)
    slice_index = index == 0 ? dims[axis] ÷ 2 + 1 : index
    
    if axis == 1
        section_data = data[slice_index, :, :]
        x_coords = coordinate_range(grid, 2)
        y_coords = coordinate_range(grid, 3)
    elseif axis == 2
        section_data = data[:, slice_index, :]
        x_coords = coordinate_range(grid, 1)
        y_coords = coordinate_range(grid, 3)
    elseif axis == 3
        section_data = data[:, :, slice_index]
        x_coords = coordinate_range(grid, 1)
        y_coords = coordinate_range(grid, 2)
    else
        throw(ArgumentError("axis must be 1, 2, or 3 for 3D grid"))
    end
    
    return x_coords, y_coords, section_data
end

# Advanced UnitGrid functionality
"""
    refine(grid::UnitGrid, factor::Int=2)

Create a refined UnitGrid with more points.
"""
function refine(grid::UnitGrid{D}, factor::Int=2) where D
    new_dims = ntuple(i -> (grid.dims[i] - 1) * factor + 1, D)
    UnitGrid{D}(new_dims)
end

"""
    coarsen(grid::UnitGrid, factor::Int=2)

Create a coarsened UnitGrid with fewer points.
"""
function coarsen(grid::UnitGrid{D}, factor::Int=2) where D
    new_dims = ntuple(i -> div(grid.dims[i] - 1, factor) + 1, D)
    UnitGrid{D}(new_dims)
end

# Periodic boundary support
"""
    periodic_distance(grid::AbstractCartesianGrid, point1, point2, periodic_axes) -> Float64

Calculate distance between two points with periodic boundary conditions.
"""
function periodic_distance(grid::G, point1, point2, periodic_axes::Vector{Bool}) where {G <: AbstractCartesianGrid}
    D = ndim(grid)
    if length(periodic_axes) != D
        throw(ArgumentError("periodic_axes must have same length as grid dimension"))
    end
    
    ext = extent(grid)
    dist_squared = 0.0
    
    for i in 1:D
        diff = point2[i] - point1[i]
        if periodic_axes[i]
            bounds = ext[i]
            domain_size = last(bounds) - first(bounds)
            # Find minimum distance considering periodicity
            diff = diff - domain_size * round(diff / domain_size)
        end
        dist_squared += diff^2
    end
    
    return sqrt(dist_squared)
end

"""
    iter_mirror_points(grid::AbstractCartesianGrid, point, periodic_axes::Vector{Bool}) -> Vector

Generate mirror points for a given point considering periodic boundaries.
"""
function iter_mirror_points(grid::G, point, periodic_axes::Vector{Bool}) where {G <: AbstractCartesianGrid}
    if length(periodic_axes) != D
        throw(ArgumentError("periodic_axes must have same length as grid dimension"))
    end
    
    ext = extent(grid)
    mirrors = [point]  # Include original point
    
    # Generate all combinations of periodic shifts
    for axis in 1:D
        if periodic_axes[axis]
            bounds = ext[axis]
            domain_size = last(bounds) - first(bounds)
            new_mirrors = typeof(mirrors)()
            
            for mirror in mirrors
                # Add shifted versions
                shifted_left = copy(mirror)
                shifted_left[axis] -= domain_size
                push!(new_mirrors, shifted_left)
                
                shifted_right = copy(mirror) 
                shifted_right[axis] += domain_size
                push!(new_mirrors, shifted_right)
                
                push!(new_mirrors, mirror)  # Keep original
            end
            mirrors = new_mirrors
        end
    end
    
    return mirrors
end